#!usr/bin/python3
#-*- coding: utf-8 -*-

import sys, os, re, time, json, random, threading, requests

from io import BytesIO
from contextlib import closing
from urllib import request


def random_ip():
    a=random.randint(1,255)
    b=random.randint(1,255)
    c=random.randint(1,255)
    d=random.randint(1,255)
    return '%s.%s.%s.%s' % (str(a), str(b), str(c), str(d))

def random_Agent():
    agents = [
    "Mozilla/5.0 (Windows NT 6.3; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/35.0.1916.153 Safari/537.36",
    "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:30.0) Gecko/20100101 Firefox/30.0",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_2) AppleWebKit/537.75.14 (KHTML, like Gecko) Version/7.0.3 Safari/537.75.14",
    "Mozilla/5.0 (compatible; MSIE 10.0; Windows NT 6.2; Win64; x64; Trident/6.0)",
    'Mozilla/5.0 (Windows; U; Windows NT 5.1; it; rv:1.8.1.11) Gecko/20071127 Firefox/2.0.0.11',
    'Opera/9.25 (Windows NT 5.1; U; en)',
    'Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1; .NET CLR 1.1.4322; .NET CLR 2.0.50727)',
    'Mozilla/5.0 (compatible; Konqueror/3.5; Linux) KHTML/3.5.5 (like Gecko) (Kubuntu)',
    'Mozilla/5.0 (X11; U; Linux i686; en-US; rv:1.8.0.12) Gecko/20070731 Ubuntu/dapper-security Firefox/1.5.0.12',
    'Lynx/2.8.5rel.1 libwww-FM/2.14 SSL-MM/1.4.1 GNUTLS/1.2.9',
    "Mozilla/5.0 (X11; Linux i686) AppleWebKit/535.7 (KHTML, like Gecko) Ubuntu/11.04 Chromium/16.0.912.77 Chrome/16.0.912.77 Safari/535.7",
    "Mozilla/5.0 (X11; Ubuntu; Linux i686; rv:10.0) Gecko/20100101 Firefox/10.0 ",
    'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:73.0) Gecko/20100101 Firefox/73.0',
    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/70.0.3538.110 Safari/537.36'
    ]
    return random.choice(agents)


def add_header(text=True):
    base = {
        "Accept-Language": "zh-CN,zh;q=0.9,en-US;q=0.8,en;q=0.7", 
        "Proxy-Connection": "keep-alive", 
        'User-Agent': random_Agent(), 
        'X-Forwarded-For':random_ip(),
    }
    if not text:
        base['Accept-Encoding'] = 'gzip, deflate, br'
    return base

def byteunit(byte):
    return '%.2fKB' % (byte / 1024.0) if byte / 1024.0 < 1024 else ('%.2fMB' % (byte / 1024.0 / 1024.0) if (byte / 1024.0 / 1024.0) < 1024 else '%.2fGB' % (byte / 1024.0 / 1024.0 / 1024.0))

def downloadfile(url, savepath='.', filename=None, params={}, headers={}, block = 1024, timeout=30, text=False):
    if not isinstance(headers, dict):
        raise BaseException("headers must be dict: %s" % headers)
    for k, v in add_header(text).items():
        headers[k] = v
    os.makedirs(savepath, exist_ok = True)
    size = 0
    if filename:
        local_filename = '%s/%s' % (savepath, filename)
    else:
        local_filename = '%s/%s' % (savepath, url.rsplit('/', 1)[1])
    tmp_filename = '%s.tmp' % local_filename
    if os.path.exists(tmp_filename):
        with open(tmp_filename, 'rb') as tmp_f:
            size = int(tmp_f.read())
    else:
        open(tmp_filename, 'wb').close()
    try:
        headers['Range'] = 'bytes=0-4'
        rep = requests.head(url, headers = headers)
        print('[+]Http Status: %s, %s' % (rep.status_code, rep.reason))
        if rep.status_code not in [200, 206]:
            raise BaseException('[+]Http Status: %s, %s' % (rep.status_code, rep.reason))
        crange = rep.headers.get('content-range', '')
        print('[DEBUG]header: %s' % rep.headers)
        # 文件大小
        rex = re.match(r'^bytes 0-4/(\d+)$', crange)
        if rex:
            total = int(rex.group(1))
            # 已下载文件大小
            headers['Range'] = 'bytes=%d-' % (size,)
        else:
            total = int(rep.headers.get('content-length', 0))
            if os.path.exists(local_filename): os.remove(local_filename)
            file_ = open(local_filename, 'ab')
            file_.close()
        print('[+]total: %s, size: %s' % (byteunit(total), byteunit(size)))
        if total <= 0: raise BaseException('total is Zero')
        rep = requests.get(url, allow_redirects=True, timeout=(30,timeout), headers=headers, params=params, stream=True)
        start_t = time.time()
        with open(local_filename, 'ab') as dst:
            dst.seek(size)
            for chunk in rep.iter_content(chunk_size = block):
                if chunk:
                    dst.write(chunk)
                    size += len(chunk)
                    dst.flush()
                print('\r%3.2f%% in %s' % (size / total * 100.0, byteunit(total)), end='')
            spend = int(time.time() - start_t)
            speed = int(size / 1024 / (spend if spend > 0 else 1))
            os.remove(tmp_filename)
            print('\nDownload finish! spend: %ds, speed: %sKB/s' % (spend, speed))
    except BaseException as e:
        with open(tmp_filename, 'wb') as tmp_f:
            tmp_f.write(('%s' % size).encode('utf-8'))
        print('Error: %s' % e)

def download_range(url, filename, start, end, params={}, block=1024, timeout = 30, text = False, mt = True, status=None):
    try:
        headers = {}
        for k, v in add_header(text).items():
            headers[k] = v
        if mt:
            headers['Range'] = 'bytes=%d-%d' % (start, end)
        thread_name = threading.current_thread().getName()
        print('[-]%s: %16s   ~%16s' % (thread_name, start, end))
        rep = requests.get(url, allow_redirects=True, headers=headers, params=params, stream=True)
#        print('%s: rsq: headers(%s), rep: headers(%s)' % (thread_name, headers, rep.headers))
        filename = '%s.%s' % (filename, start)
        if os.path.exists(filename): os.remove(filename)
        with open(filename, "ab") as fp:
#            fp.seek(start)
#            fp.write(rep.content)
            size = 0
            total = end - start
            if isinstance(status, dict):
                status[thread_name] = total / total * 100.0
            for chunk in rep.iter_content(chunk_size = block):
                fp.write(chunk)
                size += len(chunk)
                fp.flush()
                if isinstance(status, dict):
                    status[threading.current_thread().getName()] = size / total * 100.0
#            print('%s: %s, (%s, %s)%s' % (threading.current_thread().getName(), size, start, end, total))
    except BaseException as e:
        print('%s | Error: %s' % (threading.current_thread().getName(), e))
def download_mt(url, savepath, num_thread=5, filename=None):
    try:
        rep = requests.head(url, headers = {'Range': 'bytes=0-4'})
        print('[+]Http Status: %s, %s' % (rep.status_code, rep.reason))
        if rep.status_code not in [200, 206]:
            raise BaseException('[+]Http Status: %s, %s' % (rep.status_code, rep.reason))
        crange = rep.headers.get('content-range', '')
        print('[DEBUG]header: %s' % rep.headers)
        # 文件大小
        rex = re.match(r'^bytes 0-4/(\d+)$', crange)
        if rex:
            total = int(rex.group(1))
            mt = True
        else:
            total = int(rep.headers.get('content-length', 0))
            mt = False
        # 不支持多线程下载
        if not mt: num_thread = 1
        print('[+]total: %s, thread num: %s' % (byteunit(total), num_thread))
        if total <= 0: raise BaseException('total is Zero')
        os.makedirs(savepath, exist_ok = True)
        if filename:
            local_filename = '%s/%s' % (savepath, filename)
        else:
            local_filename = '%s/%s' % (savepath, url.rsplit('/', 1)[1])
        part = total // num_thread
        status = {}
        files = []
        start_t = time.time()
        for i in range(num_thread):
            start = part * i
            if i == num_thread - 1:
                end = total
            else:
                end = start + part
            if i > 0:
                start += 1
            thread_name = '%s-%s' % ('download_range', i)
            status[thread_name] = 0.0
            files.append('%s.%s' % (local_filename, start))
            t = threading.Thread(target=download_range, name=thread_name, kwargs={'url': url, 'filename': local_filename, 'start': start, 'end': end, 'mt': mt, 'status': status})
            t.start()
        t = threading.Thread(target=check_status, name='check_status', kwargs={'status': status})
        t.start()
        t.join()
        spend = int(time.time() - start_t)
        speed = int(total / 1024 / (spend if spend > 0 else 1))
        mergefiles(files, local_filename)
        print('[-]%s 下载完成' % local_filename)
        print('[*]Download finish! spend: %ds, speed: %sKB/s' % (spend, speed))
    except BaseException as e:
        print('Error: %s' % e)

def check_status(status):
    while True:
        flag = True
        process = ''
        for k, v in status.items():
            process += '%8.2f%%' % v
            if v < 100: flag = False
        print('[~]status: %s' % process, end = '\r')
        if flag: break
    print('')

def mergefiles(files, filename):
    size = len(files)
    files = [item for item in files if os.path.exists(item)]
    if size > len(files):
        raise BaseException('download file not complete')
    if os.path.exists(filename): os.remove(filename)
    with open(filename, 'ab') as dst:
        for item in files:
            with open(item, 'rb') as src:
                dst.write(src.read())
            os.remove(item)

if __name__ == '__main__':
    help_msg = '''\
        Usage:python3 %s -u url -p path
        -u: download url
        -p: savepath
        ----------------------------------------
        require requests
        pip3 install requests -i https://pypi.tuna.tsinghua.edu.cn/simple
        '''    % sys.argv[0]
    if len(sys.argv) < 2:
        print(help_msg)
    else:
        if '-m' in sys.argv:
            if '-u' in sys.argv:
                url = sys.argv[sys.argv.index('-u') + 1]
                if '-p' not in sys.argv:
                    path = url.rsplit('/', 1)[1]
                    path = path.rsplit('.', 1)[0] if '.' in path else path
                else:
                    path = sys.argv[sys.argv.index('-p') + 1]
                if '-n' in sys.argv:
                    th_num = int(sys.argv[sys.argv.index('-n') + 1])
                else:
                    th_num = 5
                if '-f' in sys.argv:
                    filename = sys.argv[sys.argv.index('-f') + 1]
                else:
                    filename = None
                download_mt(url, path, th_num, filename)
            exit(0)
        if '-u' in sys.argv:
            url = sys.argv[sys.argv.index('-u') + 1]
            if '-p' not in sys.argv:
                path = url.rsplit('/', 1)[1]
                path = path.rsplit('.', 1)[0] if '.' in path else path
            else:
                path = sys.argv[sys.argv.index('-p') + 1]
            if '-f' in sys.argv:
                filename = sys.argv[sys.argv.index('-f') + 1]
            else:
                filename = None
            downloadfile(url, path, filename)
            exit(0)
        print(help_msg)
